from visualization import show_image_label, plot_class_ratios
from dataset import FontDataset
import torch
import numpy as np


def count_class(train_dataset):
    """统计数据集中包含的类别"""
    res = set()
    for i in range(len(train_dataset)):
        if i % 100 == 0:
            print('当前执行：第', i, '个')
        _, l = train_dataset[i]
        for x in range(288):
            for y in range(288):
                res.add(l[x][y])
    print(res)
    print(len(res))


def find_class_in_image(train_dataset, target_class, number=1):
    """查找数据集中包含某类别的图像"""
    found_indices = []  # 用于存储找到的图像索引
    for i in range(len(train_dataset)):
        _, l = train_dataset[i]
        if target_class in l:
            found_indices.append(i)
        if len(found_indices) >= number:
            break  # 找到 3 个后退出
    print("最终找到的图像索引列表:", found_indices)
    print(len(found_indices))
    for index in found_indices:
        image, label = train_dataset[index]
        show_image_label(image.squeeze(0), label)
    return found_indices


def compare_labels(dataset1, dataset2, num_samples=10):
    """
    比较两个数据集中每个索引的标签类别是否一致。

    参数：
    - dataset1: 第一个数据集
    - dataset2: 第二个数据集
    - num_samples: 取前 num_samples 个样本进行对比

    返回：
    - mismatch_count: 不匹配的样本数量
    """
    print(f'dataset1 len:{len(dataset1)}')
    print(f'dataset2 len:{len(dataset2)}')
    if len(dataset1) != len(dataset2):
        print('两个数据集大小不一致！')
        return
    mismatch_count = 0
    for i in range(len(dataset1)):
        label1 = dataset1[i][1]  # 获取第 i 个标签 (320, 320)
        label2 = dataset2[i][1]  # 获取第 i 个标签 (320, 320)

        # 获取唯一类别
        unique1 = set(torch.unique(label1).tolist())  # 转换为set
        unique2 = set(torch.unique(label2).tolist())  # 转换为set

        # 移除类别 34
        unique1.discard(34)
        unique2.discard(34)

        if unique1 == unique2:
            # print(f"样本 {i}: 类别匹配 {list(unique1)}")
            pass
        else:
            print(f"样本 {i}: 数据集1类别 {list(unique1)}, 数据集2类别 {list(unique2)}")
            mismatch_count += 1

    print(f"\n对比完成，总共有 {mismatch_count} 个样本类别不匹配")


def mask_logits(logits, valid_classes_batch):
    """
    logits: [B, num_classes, H, W] 预测的 logits
    valid_classes_batch: [B, 35] one-hot 编码，表示每个样本允许的类别
    """
    B, C, H, W = logits.shape
    device = logits.device

    # 创建 mask，初始值为 -1e9，形状 [B, C, H, W]
    mask = torch.full((B, C, H, W), -1e9, device=device)

    # 将有效类别的 mask 值设置为 0
    # valid_classes_batch 是 [B, C] 形状的 one-hot 编码，表示每个样本有效的类别
    # 每个样本的 valid_classes_batch 中，有效类别的值是 1，其他是 0。
    for i in range(B):
        valid_classes = valid_classes_batch[i]  # [C] one-hot 编码
        # 在 logits 的类别维度上，对有效类别的对应位置设置为 0
        mask[i, valid_classes == 1] = 0  # 通过 one-hot 编码找到有效类别的位置并设置为 0

    # 扩展 mask 的形状为 [B, C, H, W] 以便与 logits 进行相加
    return logits + mask


def compute_class_frequency(datasets):
    """
    计算整个数据集中每个类别的出现次数

    参数:
    dataset_loader: 数据集的 DataLoader
    num_classes: 类别数

    返回:
    class_occurrences: 每个类别的出现次数 (C,)
    """
    class_occurrences = np.zeros(35)

    for dataset in datasets:
        for _, label, _ in dataset:
            label = label.numpy()

            # 获取标签中所有唯一类别
            unique_classes = np.unique(label)
            unique_classes = unique_classes[unique_classes != 0]  # 排除类别 0

            # 遍历每个唯一类别，统计出现次数
            for cls in unique_classes:
                class_occurrences[cls] += 1

    # 计算总的类别出现次数
    total_occurrences = np.sum(class_occurrences)

    # 计算每个类别的频率
    class_frequencies = class_occurrences / total_occurrences

    print(sum(class_frequencies))

    return class_frequencies


def calculate_class_ratios(datasets):
    """
    计算类别像素占比
    :param datasets:
    :return:
    """
    total_pixels = 0
    class_pixel_counts = np.zeros(35)  # 假设共有 35 类
    for dataset in datasets:
        for img, label, classes in dataset:
            label_np = label.numpy()  # 转换为 numpy 数组

            # 统计总像素数时排除类别 0 的像素
            total_pixels += np.sum(label_np)  # 只统计非零类别的像素数
            # total_pixels += np.sum(label_np != 0)  # 只统计非零类别的像素数

            # 统计每个类别的像素数量，排除类别 0
            for class_id in range(35):  # 从 1 开始，跳过类别 0
                class_pixel_counts[class_id] += np.sum(label_np == class_id)

    # 计算每个类别的占比
    class_ratios = class_pixel_counts / total_pixels
    print(sum(class_ratios))
    return class_ratios


def custom_collate_fn(batch):
    images, labels, classes = zip(*batch)
    images = torch.stack(images, dim=0)  # [B, 1, 320, 320]
    labels = torch.stack(labels, dim=0)  # [B, 320, 320]

    num_classes = 35
    batch_size = len(classes)

    class_mask = torch.zeros((batch_size, num_classes), dtype=torch.float32)
    for i, cls_list in enumerate(classes):
        class_mask[i, torch.tensor(cls_list, dtype=torch.long)] = 1.0  # 直接填充对应类别

    return images, labels, class_mask  # 直接返回 one-hot 类别信息


if __name__ == "__main__":
    if __name__ == '__main__':
        TrainDataset1 = FontDataset(False, "/Users/zenos/Downloads/CCSSD/FZLBJW2017")
        TrainDataset2 = FontDataset(True, "/Users/zenos/Downloads/CCSSD/FZLBJW2017")

        # 计算类别像素占比
        class_pixel_ratios = calculate_class_ratios([TrainDataset1, TrainDataset2])
        # class_ratios = compute_class_frequency([TrainDataset1, TrainDataset2])

        # 输出每个类别的像素数和占比，排除类别0
        for class_id in range(1, 35):  # 从 1 开始，跳过类别 0
            # print(f"类别 {class_id}: 像素数 = {class_pixel_ratios[class_id] * 100:.2f}%, 类别数 = {class_ratios[class_id] * 100:.2f}%")
            print(f"类别 {class_id}: 像素数 = {class_pixel_ratios[class_id] * 100:.2f}%")

        print(class_pixel_ratios)

        # 绘制条形图
        # plot_class_ratios(class_pixel_ratios, class_ratios)